#include "FluidBB.h"

////////////////////////////////////////

static VxColor color;
static BOOL vector_data_mode;

typedef void (*AdvectFunction)( int N, int b, float * d, float * d0, float * u, float * v, float dt );
typedef void (*LinSolveFunction)( int N, int b, float * x, float * x0, float a, float c );
typedef void (*ProjectFunction)( int N, float * u, float * v, float * p, float * div );

/*
----------------------------------------------------------------------
free/clear/allocate simulation data
----------------------------------------------------------------------
*/


void free_data ( void )
{
	if ( u )
	{
		free ( u );
		u = NULL;
	}

	if ( v ) 
	{
		free ( v );
		v = NULL;
	}

	if ( u_prev ) 
	{
		free ( u_prev );
		u_prev = NULL;
	}

	if ( v_prev )
	{ 
		free ( v_prev );
		v_prev = NULL;
	}

	if ( dens ) 
	{
		free ( dens );
		dens = NULL;
	}

	if ( dens_prev ) 
	{
		free ( dens_prev );
		dens_prev = NULL;
	}
}

void clear_data ( void )
{
	int i, size=(N+2)*(N+2);

	for ( i=0 ; i<size ; i++ ) {
		u[i] = v[i] = u_prev[i] = v_prev[i] = dens[i] = dens_prev[i] = 0.0f;
	}
}

int allocate_data ( void )
{
	int size = (N+2)*(N+2);

	u			= (float *) malloc ( size*sizeof(float) );
	v			= (float *) malloc ( size*sizeof(float) );
	u_prev		= (float *) malloc ( size*sizeof(float) );
	v_prev		= (float *) malloc ( size*sizeof(float) );
	dens		= (float *) malloc ( size*sizeof(float) );	
	dens_prev	= (float *) malloc ( size*sizeof(float) );

	if ( !u || !v || !u_prev || !v_prev || !dens || !dens_prev ) {
		fprintf ( stderr, "cannot allocate data\n" );
		return ( 0 );
	}

	return ( 1 );
}


void add_source ( int N, float * x, float * s, float dt )
{
	int i, size=(N+2)*(N+2);
	for ( i=0 ; i<size ; i++ ) x[i] += dt*s[i];
}

void set_bnd ( int N, int b, float * x )
{
	int i;

	for ( i=1 ; i<=N ; i++ ) 
	{
		x[IX(0  ,i)] = b==1 ? -x[IX(1,i)] : x[IX(1,i)];
		x[IX(N+1,i)] = b==1 ? -x[IX(N,i)] : x[IX(N,i)];
		x[IX(i,0  )] = b==2 ? -x[IX(i,1)] : x[IX(i,1)];
		x[IX(i,N+1)] = b==2 ? -x[IX(i,N)] : x[IX(i,N)];
	}

	x[IX(0  ,0  )] = 0.5f*(x[IX(1,0  )]+x[IX(0  ,1)]);
	x[IX(0  ,N+1)] = 0.5f*(x[IX(1,N+1)]+x[IX(0  ,N)]);
	x[IX(N+1,0  )] = 0.5f*(x[IX(N,0  )]+x[IX(N+1,1)]);
	x[IX(N+1,N+1)] = 0.5f*(x[IX(N,N+1)]+x[IX(N+1,N)]);
}

void lin_solve_collision ( int N, int b, float * x, float * x0, float a, float c )
{
	int i, j, k;

	for ( k=0 ; k<20 ; k++ ) 
	{
		FOR_EACH_CELL
			x[IX(i,j)] = (x0[IX(i,j)] + a*(x[IX(i-1,j)]+x[IX(i+1,j)]+x[IX(i,j-1)]+x[IX(i,j+1)]))/c;
		END_FOR
			set_bnd ( N, b, x );
	}
}

void lin_solve_overlap ( int N, int b, float * x, float * x0, float a, float c )
{
	int i, j, k;

	for ( k=0 ; k<20 ; k++ ) 
	{
		for (i = 0; i <= (N+1); i++)
			for (j = 0; j <= (N+1); j++)
			{
				//if (i == 0)
				x[IX(i,j)] = (x0[IX(i, j)] + a * (x[IX( (i == 0 ? N+1 : i - 1), j)]+x[IX((i == (N+1) ? 0 : i + 1), j)]+x[IX(i, (j == 0 ? N+1 : j - 1))]+x[IX(i, (j == (N+1) ? 0 : j + 1))])) / c;
			}
	}
}

///////////////////////////////////////////////////////

LinSolveFunction LinSolveFct[2]={lin_solve_collision,lin_solve_overlap};

///////////////////////////////////////////////////////

void diffuse ( int N, int b, float * x, float * x0, float diff, float dt )
{
	float a=dt*diff*N*N;
	LinSolveFct[border_mode] ( N, b, x, x0, a, 1+4*a );
}

void advect_overlap ( int N, int b, float * d, float * d0, float * u, float * v, float dt )
{
	int i, j, i0, j0, i1, j1;
	float x, y, s0, t0, s1, t1, dt0;

	dt0 = dt*N;

	for (i = 0; i <= (N+1); i++)
	{
		for (j = 0; j <= (N+1); j++)
		{
			x = i-dt0*u[IX(i,j)]; y = j-dt0*v[IX(i,j)];
			if (x<0.5f) x=N+0.5f; if (x>N+0.5f) x=0.5f; i0=(int)x; i1=i0+1;
			if (y<0.5f) y=N+0.5f; if (y>N+0.5f) y=0.5f; j0=(int)y; j1=j0+1;
			s1 = x-i0; s0 = 1-s1; t1 = y-j0; t0 = 1-t1;
			d[IX(i,j)] = s0*(t0*d0[IX(i0,j0)]+t1*d0[IX(i0,j1)])+
				s1*(t0*d0[IX(i1,j0)]+t1*d0[IX(i1,j1)]);
		}
	}
}

void advect_collision ( int N, int b, float * d, float * d0, float * u, float * v, float dt )
{
	int i, j, i0, j0, i1, j1;
	float x, y, s0, t0, s1, t1, dt0;

	dt0 = dt*N;

	FOR_EACH_CELL
		x = i-dt0*u[IX(i,j)]; y = j-dt0*v[IX(i,j)];
	if (x<0.5f) x=0.5f; if (x>N+0.5f) x=N+0.5f; i0=(int)x; i1=i0+1;
	if (y<0.5f) y=0.5f; if (y>N+0.5f) y=N+0.5f; j0=(int)y; j1=j0+1;
	s1 = x-i0; s0 = 1-s1; t1 = y-j0; t0 = 1-t1;
	d[IX(i,j)] = s0*(t0*d0[IX(i0,j0)]+t1*d0[IX(i0,j1)])+
		s1*(t0*d0[IX(i1,j0)]+t1*d0[IX(i1,j1)]);
	END_FOR
		set_bnd ( N, b, d );
}

///////////////////////////////////////////////////////

AdvectFunction AdvectFct[2]={advect_collision,advect_overlap};

///////////////////////////////////////////////////////

void project_collision ( int N, float * u, float * v, float * p, float * div )
{
	int i, j;

	FOR_EACH_CELL
		div[IX(i,j)] = -0.5f*(u[IX(i+1,j)]-u[IX(i-1,j)]+v[IX(i,j+1)]-v[IX(i,j-1)])/N;
	p[IX(i,j)] = 0;
	END_FOR	
		set_bnd ( N, 0, div ); set_bnd ( N, 0, p );

	LinSolveFct[border_mode] ( N, 0, p, div, 1, 4 );

	FOR_EACH_CELL
		u[IX(i,j)] -= 0.5f*N*(p[IX(i+1,j)]-p[IX(i-1,j)]);
	v[IX(i,j)] -= 0.5f*N*(p[IX(i,j+1)]-p[IX(i,j-1)]);
	END_FOR
		set_bnd ( N, 1, u ); set_bnd ( N, 2, v );
}

void project_overlap ( int N, float * u, float * v, float * p, float * div )
{
	int i, j;

	for (i = 0; i <= (N+1); i++)
	{
		for (j = 0; j <= (N+1); j++)
		{
			div[IX(i, j)] = -0.5f * (u[IX(i == (N+1) ? 0 : i+1, j)] - u[IX(i == 0 ? N + 1 : i-1, j)]+v[IX(i, j == (N+1) ? 0 : j+1)] - v[IX(i, j == 0 ? N + 1 : j-1)]) / N;
			p[IX(i, j)] = 0;
		}
	}

	LinSolveFct[border_mode] ( N, 0, p, div, 1, 4 );

	for (i = 0; i <= (N+1); i++)
	{
		for (j = 0; j <= (N+1); j++)
		{
			u[IX(i,j)] -= 0.5f * N * (p[IX(i == (N+1) ? 0 : i + 1, j)] - p[IX(i == 0 ? N + 1 : i - 1, j)]);
			v[IX(i,j)] -= 0.5f * N * (p[IX(i, j == (N+1) ? 0 : j+1)]-p[IX(i, j == 0 ? N + 1 : j - 1)]);
		}
	}
}

///////////////////////////////////////////////////////

ProjectFunction ProjectFct[2]={project_collision,project_overlap};

///////////////////////////////////////////////////////

void dens_step ( int N, float * x, float * x0, float * u, float * v, float diff, float dt )
{
	add_source ( N, x, x0, dt );
	SWAP ( x0, x ); diffuse ( N, 0, x, x0, diff, dt );
	SWAP ( x0, x ); AdvectFct[border_mode] ( N, 0, x, x0, u, v, dt );
}

void vel_step ( int N, float * u, float * v, float * u0, float * v0, float visc, float dt, CKContext *bc )
{
	add_source ( N, u, u0, dt ); add_source ( N, v, v0, dt );
	SWAP ( u0, u ); diffuse ( N, 1, u, u0, visc, dt );
	SWAP ( v0, v ); diffuse ( N, 2, v, v0, visc, dt );
	ProjectFct[border_mode] ( N, u, v, u0, v0 );
	SWAP ( u0, u ); SWAP ( v0, v );
	AdvectFct[border_mode] ( N, 1, u, u0, u0, v0, dt ); AdvectFct[border_mode] ( N, 2, v, v0, u0, v0, dt );
	ProjectFct[border_mode] ( N, u, v, u0, v0 );
}

////////////////////////////////////////


////////////////////////////////////////
// Building block
////////////////////////////////////////

void density_to_texture(float *d, CKTexture *tex, CKContext *bc)
{
	CKDWORD* data = (CKDWORD*)tex->LockSurfacePtr();
	Vx2DVector v_tmp;
	float f_tmp;
	int len = tex->GetWidth() * tex->GetHeight();

	for (int idx = 0; idx < len; idx++)
	{
		d[idx] -= loss;
		float val = d[idx];

		if (val > 1.0f)
			val = 1.0f;
		else if (val < 0.0f)
			val = 0.0f;
		else
		{
			d[idx] -= loss;
			val = d[idx];
		}

		if (vector_data_mode == FALSE)
		{
			data[idx] = RGBAFTOCOLOR(&(color * val));//RGBAFTOCOLOR(val, val, val, 1.0f);
		}
		else
		{
			v_tmp.Set(u[idx], v[idx]);
			f_tmp = v_tmp.Magnitude();
			v_tmp.Normalize();
			data[idx] = RGBAFTOCOLOR(v_tmp.x, v_tmp.y, val, f_tmp);//RGBAFTOCOLOR(val, val, val, 1.0f);
		}
	}

	tex->ReleaseSurfacePtr();
}

int UpdateFluid(const CKBehaviorContext& behcontext)
{
	CKBehavior* beh = behcontext.Behavior;
	CKContext *bc = behcontext.Context;

	//bc->OutputToConsole("OK0");

	CKTexture *tex;

	tex = (CKTexture *)beh->GetInputParameterObject(0);
	beh->GetLocalParameterValue(0, &vector_data_mode);

	if (beh->IsInputActive(0))
	{
		beh->ActivateInput(0, FALSE);	// deactivation
		beh->ActivateInput(1, FALSE);	// deactivation
		N = (tex->GetWidth() < tex->GetHeight() ? tex->GetWidth() : tex->GetHeight()) - 2;
		//behcontext.Context->OutputToConsole("alloc");
		allocate_data();
		clear_data();
		//behcontext.Context->OutputToConsole("alloc end");
	

		/*
		XString dbg = tex->GetName();
		dbg += " w: ";
		dbg += tex->GetWidth();
		dbg += " ; h : ";
		dbg += tex->GetHeight();
		bc->OutputToConsole((CKSTRING)dbg.CStr());
		beh->GetInputParameterValue(1, &dt);
		beh->GetInputParameterValue(2, &visc);
		beh->GetInputParameterValue(3, &diff);
		beh->GetInputParameterValue(4, &force);
		beh->GetInputParameterValue(5, &source);

		dbg = "Step : ";
		dbg += dt;
		dbg += " ; Visc : ";
		dbg += visc;
		dbg += " ; Diff : ";
		dbg += diff;
		dbg += " ; Force : ";
		dbg += force;
		dbg += " ; Source : ";
		dbg += source;

		bc->OutputToConsole((CKSTRING)dbg.CStr());
		*/
	}

	if (beh->IsInputActive(1))
	{
		beh->ActivateInput(0, FALSE);	// deactivation
		beh->ActivateInput(1, FALSE);	// deactivation

		beh->GetInputParameterValue(1, &border_mode);
		beh->GetInputParameterValue(2, &dt);
		beh->GetInputParameterValue(3, &visc);
		beh->GetInputParameterValue(4, &diff);
		beh->GetInputParameterValue(5, &loss);
		
		if (vector_data_mode == FALSE)
			beh->GetInputParameterValue(6, &color);
		
		vel_step ( N, u, v, u_prev, v_prev, visc, dt, bc );		
		dens_step ( N, dens, dens_prev, u, v, diff, dt );		
		density_to_texture(dens, tex, bc);

		int i, j;

		FOR_EACH_CELL
			u_prev[IX(i,j)] = v_prev[IX(i,j)] = dens_prev[IX(i,j)] = 0.0f;
		END_FOR
	}

	beh->ActivateOutput(1);
	return (CKBR_OK);
}

CKERROR UpdateFluidCallBack(const CKBehaviorContext& behcontext)
{
	CKBehavior* beh = behcontext.Behavior;

	switch( behcontext.CallbackMessage )
	{
	case CKM_BEHAVIORRESUME:
		{
			break ;
		}
	case CKM_BEHAVIORRESET:
		{
			free_data();
			//behcontext.Context->OutputToConsole("delete");
			break ;
		}
	case CKM_BEHAVIORSETTINGSEDITED:
		{
			beh->GetLocalParameterValue(0, &vector_data_mode);
			//XString tmp = "setting : v_data = ";
			//tmp += vector_data_mode;
			if (vector_data_mode == FALSE)
			{
				beh->CreateInputParameter("Color", CKPGUID_COLOR);

				//behcontext.Context->OutputToConsole((CKSTRING)tmp.CStr());

				//color.Set(1.0f, 1.0f, 1.0f, 1.0f);
				//beh->AddInputParameter("Color", CKPGUID_COLOR, &color, sizeof(VxColor));
			}
			else
			{
				beh->RemoveSubBehaviorLink(6);	
				beh->RemoveInputParameter(6);
			}
			break ;
		}
	}

	return CKBR_OK; 
}

CKERROR CreateUpdateFluidProto ( CKBehaviorPrototype **pproto)
{
	CKBehaviorPrototype *proto = CreateCKBehaviorPrototype("UpdateFluid");
	u = v = u_prev = v_prev = dens_prev = dens = NULL;
	color.Set(1.0f, 1.0f, 1.0f, 1.0f);
	vector_data_mode = FALSE;

	if(!proto)
		return (CKERR_OUTOFMEMORY);

	proto->DeclareInput("In");
	proto->DeclareInput("LoopIn");
	proto->DeclareOutput("Out");
	proto->DeclareOutput("LoopOut");

	proto->DeclareInParameter("Texture", CKPGUID_TEXTURE, 0, 0);
	proto->DeclareInParameter("Border mode", CKPGUID_BORDERMODE, "0");
	proto->DeclareInParameter("Step", CKPGUID_FLOAT, "0.1");
	proto->DeclareInParameter("Viscosity", CKPGUID_FLOAT, "0.0");
	proto->DeclareInParameter("Diffusion", CKPGUID_FLOAT, "0.0");
	proto->DeclareInParameter("Loss", CKPGUID_FLOAT, "0.0");
	proto->DeclareInParameter("Color", CKPGUID_COLOR, &color, sizeof(VxColor));

	proto->DeclareSetting("Fluid data in pixel color (dir.x,dir.y,density,force)", CKPGUID_BOOL);

	proto->SetBehaviorCallbackFct(UpdateFluidCallBack);
	proto->SetFlags(CK_BEHAVIORPROTOTYPE_NORMAL);
	proto->SetFunction(UpdateFluid);
	//proto->SetBehaviorFlags(CKBEHAVIOR_TARGETABLE);

	*pproto = proto;


	return (CK_OK);
}

CKObjectDeclaration *FillBehaviorUpdateFluid()
{
	CKObjectDeclaration *od = CreateCKObjectDeclaration ( "UpdateFluid" );
	od->SetDescription ( "Updates the velocity and density fields of the fluid." );
	od->SetCategory ( "Fluids" );
	od->SetType ( CKDLL_BEHAVIORPROTOTYPE );
	od->SetGuid ( UPDATEFLUID_GUID );
	od->SetAuthorGuid ( AUTHOR_GUID );
	od->SetAuthorName ( "Jonathan DERROUGH" );
	od->SetVersion ( 0x00000001 );
	od->SetCreationFunction ( CreateUpdateFluidProto );
	//od->SetCompatibleClassId ( CKCID_3DENTITY );
	return od;
}
